import copy
import logging

from myapp.configs.auto_learning import AUTO_LEARNING_CFGS
from myapp.configs.deploy_service import CUSTOM_CFGS, SPECIFICATIONS
from myapp.configs.training import COMPRESS_CONFIGS, RESOURCE_CONFIG, TRAIN_CONFIGS, RLHF_TRAIN_CONFIGS


# 获取服务部署资源配置信息
def get_deploy_resource_cfgs():
    return SPECIFICATIONS.to_dict()


def get_train_resource_cfgs():
    return RESOURCE_CONFIG.to_dict()


def is_multi_node(resource_cfg):
    """ " 判断规格资源是否需要多机多卡"""
    node_num = resource_cfg.get('node_num', 1)
    if node_num > 1:
        return True
    else:
        return False


def get_resource(flavor_id):
    for resource in RESOURCE_CONFIG.to_dict().values():
        if resource.get(flavor_id):
            return resource.get(flavor_id)
    logging.warning(f'[get_resource] {flavor_id} not found in TRAINING:resource')
    return None


def get_deploy_resource(flavor_id):
    for resource in SPECIFICATIONS.to_dict().values():
        if resource.get(flavor_id):
            return resource.get(flavor_id)
    logging.warning(f'[get_deploy_resource] {flavor_id} not found in DEPLOYMENT:specifications')
    return None


def get_train_resources(
    task_type, lora=False, flavor_key='exclusive_flavor_id', region_key='default'
):
    """通过task_type，获取训练/微调任务对应的可运行的规格配置，公共资源、专属资源均可获取"""
    resources = {}

    try:
        train_cfg = TRAIN_CONFIGS.get(task_type)
        logging.info(
            f'[get_train_resources] task_type: {task_type}; '
            f'flavor_key: {flavor_key}; train_cfg: {train_cfg}'
        )

        for k, v in train_cfg.items():
            resources[k] = []
            if lora is True:
                exclusive_flavor_ids = (
                    v.get('lora', {}).get('resources', {}).get(region_key, {}).get(flavor_key, [])
                )
            else:
                exclusive_flavor_ids = (
                    v.get('resources', {}).get(region_key, {}).get(flavor_key, [])
                )

            for flavor_id in exclusive_flavor_ids:
                config = get_resource(flavor_id)
                item = copy.deepcopy(config)
                item['flavor_id'] = flavor_id
                resources[k].append(item)
    except Exception as e:
        logging.exception(e)

    return resources

def get_rlhf_train_resources(
    task_type, train_method, lora=False, flavor_key='exclusive_flavor_id', region_key='default'
):
    """通过task_type，获取训练/微调任务对应的可运行的规格配置，公共资源、专属资源均可获取"""
    resources = {}
    lora_method = "lora" if lora else "full" 
    try:
        train_cfg = RLHF_TRAIN_CONFIGS.get(task_type)
        logging.info(
            f'[get_train_resources] task_type: {task_type}; '
            f'flavor_key: {flavor_key}; train_cfg: {train_cfg}'
        )
        for k, v in train_cfg.items():
            resources[k] = []
            exclusive_flavor_ids = (
                    v.get(train_method, {}).get(lora_method, {}).get('resources', {}).get( \
                        region_key, {}).get(flavor_key, [])
                )
            for flavor_id in exclusive_flavor_ids:
                config = get_resource(flavor_id)
                item = copy.deepcopy(config)
                item['flavor_id'] = flavor_id
                resources[k].append(item)
    except Exception as e:
        logging.exception(e)

    return resources


def get_compress_resources(task_type, flavor_key='exclusive_flavor_id', region_key='default'):
    """通过task_type，获取训练/微调任务对应的可运行的规格配置"""
    resources = {}

    try:
        train_cfg = COMPRESS_CONFIGS.get(task_type)
        logging.info(
            f'[get_train_resources] task_type: {task_type}; '
            f'flavor_key: {flavor_key}; train_cfg: {train_cfg}'
        )

        for k, v in train_cfg.items():
            resources[k] = []
            flavor_ids = v.get('resources', {}).get(region_key, {}).get(flavor_key, [])

            for flavor_id in flavor_ids:
                config = get_resource(flavor_id)
                item = copy.deepcopy(config)
                item['flavor_id'] = flavor_id
                resources[k].append(item)
    except Exception as e:
        logging.exception(e)

    return resources


def get_deploy_resources(task_type, flavor_key='specifications', region_key='default'):
    """通过task_type、region_key获取部署任务对应的可运行的规格配置"""
    resources = {}

    try:
        deploy_cfg = AUTO_LEARNING_CFGS.get(task_type)
        logging.info(
            f'[get_deploy_resources] task_type: {task_type}; '
            f'flavor_key: {flavor_key}; deploy_cfg: {deploy_cfg}'
        )

        for k, _ in deploy_cfg.items():
            if k in ['GPU', 'NPU', 'DCU', 'CPU']:
                resources[k] = []
                flavors = (
                    deploy_cfg[k].get('resources', {}).get(region_key, {}).get(flavor_key, [])
                )

                if isinstance(flavors, str):
                    flavors = [flavors]

                for flavor_id in flavors:
                    config = get_deploy_resource(flavor_id=flavor_id)
                    item = copy.deepcopy(config)
                    item['flavor_id'] = flavor_id
                    resources[k].append(item)
    except Exception as e:
        logging.exception(e)

    return resources


def get_custom_resources(region_key='default', task='train'):
    """通过region_key获取自定义任务的部署任务的规格配置"""
    custom_resources = {}
    try:
        if task == 'train':
            resource_cfgs = get_train_resource_cfgs()
        else:
            resource_cfgs = get_deploy_resource_cfgs()

        custom_resources_key = CUSTOM_CFGS.get(region_key, {}).get(task, {})

        for key in custom_resources_key.keys():
            custom_resources[key] = {}

            for resource in custom_resources_key[key]:
                if key in resource_cfgs.keys():
                    if resource in resource_cfgs[key].keys():
                        custom_resources[key][resource] = resource_cfgs[key][resource]

    except Exception as e:
        logging.exception(e)

    return custom_resources


# 返回所有的配置，命中的置为可用，其它置为不可用
def get_all_resource(flavor_id):
    result = []
    for _, v in RESOURCE_CONFIG.to_dict().items():
        for kk, vv in v.items():
            item = copy.deepcopy(vv)
            if kk == flavor_id:
                item['valid'] = True
            else:
                item['valid'] = False
            item['flavor_id'] = kk
            price = vv.get('price', 0)
            item['price'] = price / 10000
            item['price_desc'] = f'{price/10000}元/分钟'
            result.append(item)
    return result


def get_resource_type(flavor_id):
    res = RESOURCE_CONFIG.to_dict()
    for k in res.keys():
        v = res.get(k)
        if v:
            val = v.get(flavor_id)
            if val:
                return k
    return 'GPU'


def postprocess_public_resource_pool(public_resource_pool):
    """对公共资源配置进行后处理：如添加price_desc、level"""
    for k in public_resource_pool.keys():
        i = 0
        for item in public_resource_pool[k]:
            price = item.get('price', 0)
            item['price'] = price / 10000
            item['price_desc'] = f'{price / 10000}元/分钟'
            if i == 0:
                item['level'] = 'lower'
            elif i == 1:
                item['level'] = 'normal'
            else:
                item['level'] = 'high'
            i += 1
    return public_resource_pool
